# -*- coding: utf-8 -*-
"""
Created on Sun Oct 11 10:09:35 2020
用于双量子点量子模型
单量子比特的归零操作
多个初始态对多个目标态的归零保真度
选用两种策略：
先采用 要么最优，要么最差
如果效果不好就选第二种
动作要么选最优的，要么选次优的

@author: Waikikilick
"""

import numpy as np
from scipy.linalg import expm
from time import *
import multiprocessing as mp
import copy
np.random.seed(1)
T = 2*np.pi
dt = np.pi/5
step_max = T/dt
sx = np.mat([[0, 1], [1, 0]], dtype=complex) 
sz = np.mat([[1, 0], [0, -1]], dtype=complex)
action_space = np.array([0,1,2,3])#,5,6,7,8])

theta_num = 6 #除了 0 和 Pi 两个点之外，点的数量
varphi_num = 21#varphi 角度一圈上的点数
#总点数为 theta_num * varphi_num + 2(布洛赫球两极)

theta = np.linspace(0,np.pi,theta_num+2,endpoint=True) 
varphi = np.linspace(0,np.pi*2,varphi_num,endpoint=False) 

def psi_set():
    psi_set = []
    for ii in range(1,theta_num+1):
        for jj in range(varphi_num):
            psi_set.append(np.mat([[np.cos(theta[ii]/2)],[np.sin(theta[ii]/2)*(np.cos(varphi[jj])+np.sin(varphi[jj])*(0+1j))]]))
    psi_set.append(np.mat([[1], [0]], dtype=complex))
    psi_set.append(np.mat([[0], [1]], dtype=complex))
    return psi_set

target_set = psi_set()
init_set = psi_set()
#----------------------------------------------------------------------------------------------------
#动作直接选最优的
def step(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        H = float(action_space[action])* sz + 1 * sx
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_

#动作选最优的，或者最差的
def step1(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        H = float(action_space[action])* sz + 1 * sx
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
    
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        # del action_list[fid_list.index(max(fid_list))]
        # del psi_list[fid_list.index(max(fid_list))]
        # del fid_list[fid_list.index(max(fid_list))]
        
        # best_action = fid_list.index(max(fid_list))
        # best_fid = max(fid_list)
        
        best_action = fid_list.index(min(fid_list))
        best_fid = min(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_

#动作选最优的，或者次优的
def step2(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        H = float(action_space[action])* sz + 1 * sx
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
        
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = fid_list[best_action]
        
    else:
        psi_list_ = copy.deepcopy(psi_list)
        fid_list_ = copy.deepcopy(fid_list)
        
        del psi_list_[fid_list_.index(max(fid_list_))]
        del fid_list_[fid_list_.index(max(fid_list_))]
        
        best_action = fid_list.index(max(fid_list_))
        
        best_fid = max(fid_list_)
        
    psi_ = psi_list[best_action]
    
    return best_action, best_fid, psi_
#---------------------------------------------------------------------------------
#将测试集的保真度从小到大排列出来，来展示保真度分布
def sort_fid(test_fidelity_list):
    sort_fid = []
    for i in range (test_fidelity_list.shape[0]):
        b = test_fidelity_list[i,:]
        sort_fid  = np.append(sort_fid,b)
    sort_fid.sort()
    return sort_fid
#--------------------------------------------------------------------------------


def job(target_psi):
    fid_list, time_list = [], []
    for psi1 in init_set:
        
        psi = psi1
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        fid_max = 0
        fid_max1 = 0
        fid_max2 = 0
        fid_max0 = 0
        
        
        step_n = 0
        start_time = time()
        while True:
            action, F, psi_ = step1(psi,target_psi,F)
            fid_max1 = max(F,fid_max1)
            psi = psi_
            step_n += 1
            if fid_max1>0.999 or step_n>step_max:
                break
        end_time = time()
        time_list.append(end_time-start_time)
        
        
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        start_time = time()
        while True:
            action, F, psi_ = step2(psi,target_psi,F)
            fid_max2 = max(F,fid_max2)
            psi = psi_
            step_n += 1
            if fid_max2>0.999 or step_n>step_max:
                break 
        end_time = time()
        time_list.append(end_time-start_time)
        fid_max = max(fid_max1,fid_max2)
        
        
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        start_time = time()
        while True:
            action, F, psi_ = step(psi,target_psi,F)
            fid_max0 = max(F,fid_max0)
            psi = psi_
            step_n += 1
            if fid_max0>0.999 or step_n>step_max:
                break   
        end_time = time()
        time_list.append(end_time-start_time)
        fid_max = max(fid_max,fid_max0)
        
        fid_list.append(fid_max)
    return  (np.mean(time_list),np.mean(fid_list))

def multicore():
    pool = mp.Pool()
    data = pool.map(job, target_set)
    return data
    

if __name__ == '__main__':
    # print(target_set)
    time1 = time()
    data = multicore()
    data.sort()
    # print(data)
    time2 = time()
    # print(np.mean(F_list))
    print('time_cost is: ',time2-time1)
    print(data)
    
    #F_list.sort()# jiang liebiao an shunxu pailie
    
# [(0.018929827958345413, 0.9986917328946125), (0.019288037593166035, 0.9985916617642765), (0.019474629312753677, 0.998526461278199), (0.019728801523645718, 0.9915930720169693), (0.019886753832300503, 0.9982392983276178), (0.0200070608407259, 0.9912947461197672), (0.02006417140364647, 0.997734366702414), (0.020103797937432926, 0.9937820420466081), (0.020162111148238182, 0.9978367009373044), (0.020331252987186115, 0.9930110516765831), (0.020352350547909737, 0.9989008835000456), (0.020452169701457024, 0.9971294953234213), (0.020508234078685444, 0.9919871130594137), (0.020512747888763744, 0.9937676273926017), (0.020512815564870834, 0.9948242129386097), (0.02054108741382758, 0.9948781930796775), (0.020565953726569813, 0.9959810391044307), (0.02056928041080634, 0.9920853583168253), (0.020613736162583034, 0.9955907074419874), (0.020632240921258926, 0.9957871392504156), (0.020645547658205032, 0.9720411240313194), (0.0206985039015611, 0.9954707593578165), (0.020700011402368546, 0.9968625424072268), (0.02073112813134988, 0.9948922528286124), (0.02073770264784495, 0.9906551527586371), (0.020765370378891628, 0.984862426083432), (0.02081303546826045, 0.9478389565791276), (0.020825161909063656, 0.991430297399177), (0.02084280674656232, 0.9286408850871837), (0.020897651712099712, 0.9978374519026476), (0.020920267949501675, 0.8319793129680799), (0.02096577112873395, 0.9829579366661805), (0.02098373572031657, 0.9879055548526442), (0.021043828999002773, 0.9970012758011078), (0.021070312708616257, 0.9712503451885746), (0.021078009779254597, 0.9730592111688531), (0.021080744763215382, 0.9914324363525823), (0.021087707951664925, 0.8656387227181339), (0.02111741341650486, 0.9857090542253422), (0.021120257675647736, 0.9946062397678472), (0.02118421842654546, 0.9942918752211851), (0.021229954436421394, 0.9768194861334132), (0.021249710271755855, 0.9948096623929437), (0.021260630960265797, 0.9962209739970853), (0.021263161053260166, 0.9509199487553386), (0.02130167745053768, 0.9912574490951213), (0.021307678893208504, 0.9436409673870634), (0.02130981038014094, 0.7957292923461645), (0.021313559884826343, 0.9903879222325596), (0.021314232299725216, 0.9886491644509754), (0.02131671831011772, 0.9889882720652874), (0.021316854283213615, 0.9733769694744154), (0.02131989598274231, 0.9889886893246), (0.021324428419272106, 0.9957995118939653), (0.021327079584201176, 0.9900675390465744), (0.021341463550925255, 0.9928267963250814), (0.021342075740297634, 0.9749152340266849), (0.02138678915798664, 0.9878645782396962), (0.02140483135978381, 0.9663421131510306), (0.021406643092632294, 0.9806549800856832), (0.02140938490629196, 0.9345583146736149), (0.02141251725455125, 0.9903229788360068), (0.02141635740796725, 0.9916937244569652), (0.021457795674602192, 0.9266097510553934), (0.02146207292874654, 0.9543483215365164), (0.021475377803047497, 0.9387908383987107), (0.02148904651403427, 0.9958746583898326), (0.02150091404716174, 0.9955922301094524), (0.021509839221835136, 0.9127178165464349), (0.021512522051731747, 0.9783443778942091), (0.021597659215331078, 0.9990775556256382), (0.021609360352158546, 0.9899538718298324), (0.021625996256868046, 0.9564341135404315), (0.021634213626384735, 0.9943252468968324), (0.021650441611806553, 0.9946591102148332), (0.021670759345094364, 0.9887744348909606), (0.021676117554306984, 0.886104620700943), (0.021677072470386822, 0.9879469284783713), (0.02170419382552306, 0.9782487018367949), (0.021710964540640514, 0.9901710237625261), (0.021740172679225605, 0.9863614214405408), (0.02174314173559348, 0.9166525420148361), (0.02175101637840271, 0.9915895100559452), (0.02179363990823428, 0.9402266190422934), (0.021802054718136787, 0.9867668208585385), (0.02180246201654275, 0.9919703975837273), (0.021979437520106632, 0.9211748737008398), (0.021991103887557983, 0.9933963452625771), (0.022008380542198818, 0.9931742226899848), (0.022012184063593548, 0.9905488513787031), (0.022147303447127342, 0.9875100370670282), (0.022179216767350834, 0.9951351255483176), (0.022210024918119114, 0.8028664867284254), (0.022211817403634388, 0.9425036504128506), (0.022265940283735592, 0.9751390649984744), (0.02234664186835289, 0.9831897649771133), (0.02236851987739404, 0.9145908143554086), (0.022396770616372425, 0.9919048768732022), (0.02240098702410857, 0.9948534843970581), (0.022431867197155952, 0.9866778069560986), (0.022519270579020183, 0.9970097078912608), (0.022590930263201397, 0.9880732980256026), (0.02264893737932046, 0.9741638695649669), (0.02266000583767891, 0.9135883702283729), (0.022722665841380756, 0.9952034312007518), (0.02280457317829132, 0.9752240182436842), (0.022810791929562885, 0.9979634769664035), (0.02282944756249587, 0.9961232975521294), (0.022921300182739895, 0.9916303556264001), (0.022962032506863277, 0.8614099677128284), (0.022967544694741566, 0.9191505208419719), (0.023179218793908756, 0.9832462282117398), (0.02327379584312439, 0.9859197186188753), (0.023346584290266037, 0.9949054767092369), (0.023382363840937614, 0.9797826837738739), (0.02340182972451051, 0.9692857475840906), (0.023421211789051693, 0.9909465281978896), (0.023428072532018025, 0.9924676880094832), (0.023456259941061337, 0.9964424528187474), (0.023476652180155117, 0.9723382875522987), (0.02353738807141781, 0.9751195495226357), (0.023593631262580555, 0.990013665733689), (0.023708293214440346, 0.9747268921739078), (0.023839217921098072, 0.9541288167805854), (0.02400428553422292, 0.9523262300317582), (0.024619845673441887, 0.9679130999500374), (0.02466760389506817, 0.8050867145043826), (0.02534639711181323, 0.9892062839412421)]
